#!/usr/bin/env python
import imageTools.ITKUtils.io as io
from imageTools.ITKUtils.itkmath import sumImages
import itk
import sys
import imageTools.browse.a3d as a3d
import cPickle
import numpy as na
import os
import random
class BoneMask(object):
    def __init__(self, img=None, imgfile=None):
        self.img = img
        self.imgfile = imgfile
class Segment(object):
    def __init__(self,imgfile, pointsfile, thresh, refimg = None ):
        self.imgfile = imgfile
        self.pointsfile = pointsfile
        self.thresh = thresh
        self.nimg = None
        self.filter = None
        self.points = None
        self.normalize = False
        self.filter = itk.ConnectedThresholdImageFilter.IF3IF3.New()
        self.vesselSeeds = None
        self.cumImg = None
        self.numSegs = 1
        self.numSeeds = 1
    def setNumSeeds(self, numSeeds):
        """set the number of seeds to use for the segmentation"""
        self.numSeeds = numSeeds
    def getImage(self):
        img = io.readImage(self.imgfile,imgMode='float')
        lr = img.GetLargestPossibleRegion()
        self.dim = lr.GetSize()
        if( self.normalize ):
            nfilter = itk.NormalizeImageFilter.IF3IF3.New()
            nfilter.SetInput(img)
            self.nimg = nfilter.GetOutput()
            self.nimg.Update()
        else:
            self.nimg = img
        self.filter.SetInput(self.nimg)
    def process(self):
        """run through the segmentation"""
        if( self.nimg == None ):
            self.getImage()
        if( self.points == None ):
            self.getCriticalPoints()
        self.setVesselSeeds()
        self.filter.SetLower(self.thresh)
        self.mask = self.filter.GetOutput()
        self.mask.Update()
    def bootSegmentationProcess(self, numBoots = 1):
        for i in xrange(numBoots):
            self.process()
            self.addResult()
    def getCriticalPoints(self):
        """read the specified critical points file"""
        fo = open(self.pointsfile,'rb')
        self.points = cPickle.load(fo)
        fo.close()
    def addSeedsCP(self):
        self.seeds = []
        for d in self.points:
            if( d.getLabel() == 'vessel' ):
                if( random.random() >= 0.0 ):
                    crd = (int(d[0]),int(d[1]),int(d[2]))
                    value = self.nimg.GetPixel(crd)
                    self.seeds.append(d)
                    self.filter.AddSeed(crd)
                    print crd, value
        #print self.seeds
    def setVesselSeeds(self):
        if not self.vesselSeeds :
            self.vesselSeeds = []
            for d in self.points:
                if( d.getLabel() == 'vessel' ):
                    self.vesselSeeds.append((int(d[0]),int(d[1]),int(d[2])))
                    
        print "selecting %d seeds from %d vessel seeds"%(self.numSeeds,len(self.vesselSeeds))
        random.shuffle(self.vesselSeeds)
        self.seeds = self.vesselSeeds[0:self.numSeeds]
        for seed in self.seeds:
            self.filter.AddSeed(seed)
            print seed, self.nimg.GetPixel(seed)
    def addResult(self):
        """Add the current segmentation result to the cummulative image"""
        if( not self.cumImg ):
            self.cumImg = self.mask
        else:
            self.cumImg = sumImages([self.cumImg, self.mask])
        
    def viewImage(self, whichImg = "mask"):
        if( whichImg == 'mask'):
            data = itk.PyBuffer.IF3.GetArrayFromImage(self.mask).astype(na.float)
        elif( whichImg == 'cum'):
            data = itk.PyBuffer.IF3.GetArrayFromImage(self.cumImg).astype(na.float)
        else:
            data = itk.PyBuffer.IF3.GetArrayFromImage(self.nimg).astype(na.float)
        a3d.call(data=[data])

    def evaluatePoints(self, pfile = None ):
        if( pfile ):
            fo = open(pfile,'rb')
            points = cPickle.load(fo)
            fo.close()
        else:
            points = self.points
        for p in points:
            print p,"nimg",self.nimg.GetPixel((int(p[0]),int(p[1]),int(p[2]))),\
                  "mask",self.mask.GetPixel((int(p[0]),int(p[1]),int(p[2])))
    
def main():
    seg = Segment(sys.argv[1], sys.argv[3], float(sys.argv[2]))
    
    seg.setNumSeeds(int(sys.argv[4]))
    seg.bootSegmentationProcess(numBoots = int(sys.argv[5]))
    
    # create a filename for saving cumulative image
    temp = os.path.splitext(sys.argv[1])
    fname = temp[0]+"seeds_%03d_boots_%03d_thresh_%03d"+temp[1]
    io.writeImage(seg.cumImg,fname=fname, outMode="uchar" )
    
    seg.viewImage("cum")
    
if __name__ == '__main__':
    main()
    

